import torchvision
from torch.utils.data import DataLoader
from torchvision.transforms import transforms

test_set = torchvision.datasets.CIFAR10(root="./dataset", train=False, transform=transforms.ToTensor())

# 使用Dataloader
# dataset:要操作的数据集
# batch_size:划分的大小
# shuffle:数据集是否重新洗牌
# drop_last: 最后划分后，不满足batch_size大小的数据集是否丢弃
img_dataloader = DataLoader(dataset=test_set, batch_size=128, shuffle=True, drop_last=False)

# 查看dataset中第一个数据
img, target = test_set[0]
print("dataset - img.shape:%s, target:%s" % (img.shape, target))

from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter("dataloader的使用logs")

count = 0
for datas,targets in img_dataloader:
    # print("dataloader - datas.shape:%s, targets:%s" % (datas.shape, targets))
    writer.add_images("test dataloader1", datas, count)
    count += 1

writer.close()